import copy
import torch
from Aggregation.BaseAggregation import BaseAggregation


class FedAvg(BaseAggregation):
    """
    联邦平均
    """

    def __init__(self, config, train_dataset, test_dataset, user_groups, traindata_cls_counts):
        super(FedAvg, self).__init__(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)

    def update_weights(self):
        """
        Returns the average of the weights.
        """
        w = self.net
        w_avg = copy.deepcopy(w[0].state_dict())
        for key in w_avg.keys():
            for i in range(1, len(w)):
                w_avg[key] += w[i].state_dict()[key]
            if 'num_batches_tracked' in key:
                w_avg[key] = w_avg[key].true_divide(len(w))
            else:
                w_avg[key] = torch.div(w_avg[key], len(w))

        self.global_model.load_state_dict(w_avg)
        self.distribute_model()

        return w_avg


